from fastapi import FastAPI, UploadFile, File, Request
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
import os
from datetime import datetime
from predicate import predict_image
app = FastAPI()
app.mount("/image_data", StaticFiles(directory="image_data"), name="image_data")

# 设置模板和静态文件
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")


# 创建必要文件夹
os.makedirs("image_data", exist_ok=True)
os.makedirs("result_data", exist_ok=True)

@app.post("/upload/")
async def upload_image(file: UploadFile = File(...)):
    # 生成时间格式文件名
    now = datetime.now()
    filename = now.strftime("%Y-%m-%d-%H-%M-%S") + ".png"
    filepath = os.path.join("image_data", filename)

    # 保存图片
    with open(filepath, "wb") as buffer:
        buffer.write(await file.read())

    # 预测图片
    pred_class, confidence = predict_image(filepath)

    # 写入结果文件
    result_line = f"{filename} {pred_class} {confidence}\n"
    with open("result_data/res.txt", "a") as f:
        f.write(result_line)

    return {
        "filename": filename,
        "prediction": pred_class,
        "confidence": f"{confidence*100:.2f}%",
        "image_url":f"/image_data/{filename}"
    }

@app.get("/")
async def main(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})
